/* MeanShift.cpp 
-----------------------------------------
By: 	Wei YE ( ye0003ei@e.ntu.edu.sg )

Date:	10/Oct/2012
-----------------------------------------
*/
#include "MeanShift.h"

void MeanShift::LoadPoints(float* data, int n_points, int n_dim)
{
	n_ = n_points;
	d_ = n_dim;
	if(haspoints_){
		delete []points_;
		delete []convpts_;
	}

	points_ = new MS_Point[n_];
	convpts_ = new float[n_*d_];
	haspoints_ = true;

	for(int p = 0; p < n_; p++){
		points_[p] = MS_Point( &data[p*d_], d_ );
	}
}

void MeanShift::NonGridMsFilter()
{
	// General MeanShift ... no Grid information is considered.
	// not used in this assignment ... so~~
}

// ---------------------------------------------------------------//
//       Mean Shift procedure proposed by Comaniciu et.al.        //
//----------------------------------------------------------------//
void MeanShift::GridMsFilter(int height, int width, int sigmas, int sigmar)
{
	//ofstream log;
	//log.open("log.txt", ios::out);

	if( haspoints_ = false){
		//	log<<"Load Points first!"<<endl;
		return;
	}
	if( height* width != n_){
		//	log<<"Grid do not match!"<<endl;
		return;
	}

	// index the data in the 3d buckets (x, y, L)
	int i,j;
	int* buckets;
	int* slist;
	slist = new int[n_];
	int bucNeigh[27];

	float sMins; // just for L
	float sMaxs[3]; // for all
	sMaxs[0] = (float)height/sigmas;
	sMaxs[1] = (float)width/sigmas;
	sMins = sMaxs[2] = points_[0].data_[2];

	float cval;
	for(i=0; i<n_; i++)
	{
		cval = points_[i].data_[2];
		if (cval < sMins)
			sMins = cval;
		else if (cval > sMaxs[2])
			sMaxs[2] = cval;
	}

	int nBuck1, nBuck2, nBuck3;
	int cBuck1, cBuck2, cBuck3, cBuck;
	nBuck1 = (int) (sMaxs[0] + 3);
	nBuck2 = (int) (sMaxs[1] + 3);
	nBuck3 = (int) (sMaxs[2] - sMins + 3);
	buckets = new int[nBuck1*nBuck2*nBuck3];
	for(i=0; i<(nBuck1*nBuck2*nBuck3); i++)
		buckets[i] = -1;

	for(i=0; i<n_; i++)
	{
		// find bucket for current data and add it to the list
		cBuck1 = (int) points_[i].data_[0] + 1;
		cBuck2 = (int) points_[i].data_[1] + 1;
		cBuck3 = (int) (points_[i].data_[2] - sMins) + 1;
		cBuck = cBuck1 + nBuck1*(cBuck2 + nBuck2*cBuck3);

		slist[i] = buckets[cBuck];
		buckets[cBuck] = i;
	}

	// init bucNeigh
	i = 0;
	for (cBuck1=-1; cBuck1<=1; cBuck1++)
	{
		for (cBuck2=-1; cBuck2<=1; cBuck2++)
		{
			for (cBuck3=-1; cBuck3<=1; cBuck3++)
			{
				bucNeigh[i++] = cBuck1 + nBuck1*(cBuck2 + nBuck2*cBuck3);
			}
		}
	}
	double hiLTr = 80.0/sigmar;

	int idxd, iterationCount, pt_count;
	float el, diff;					// difference in each component, total difference
	float* wc = new float[d_];
	MS_Point Wc = MS_Point(wc, d_); // (kernel) Window Center
	float* mt = new float[d_];
	MS_Point Mt = MS_Point(mt, d_); // creat a Mean shifT vector 
	float MtAbs;                    // abs of Mean shifT vector

	for( i = 0; i < n_; i++)
	{
		pt_count = 0;
		points_[i].cloneTo(Wc); // move the window center to current data point
		Mt.zero();

		cBuck1 = (int) Wc.data_[0] + 1;
		cBuck2 = (int) Wc.data_[1] + 1;
		cBuck3 = (int)(Wc.data_[2] - sMins) + 1;
		cBuck = cBuck1 + nBuck1*(cBuck2 + nBuck2*cBuck3);
		for (j = 0; j < 27; j++)
		{
			idxd = buckets[cBuck+bucNeigh[j]];
			while (idxd >= 0)
			{
				// determine if inside search window
				el = points_[idxd].data_[0] - Wc.data_[0];
				diff = el*el;
				el = points_[idxd].data_[1] - Wc.data_[1];
				diff += el*el;

				if (diff < 1.0)
				{
					el = points_[idxd].data_[2] - Wc.data_[2];
					if (Wc.data_[2] > hiLTr)
						diff = 4*el*el;
					else
						diff = el*el;

					if (d_ == 5)
					{
						el = points_[idxd].data_[3] - Wc.data_[3];
						diff += el*el;
						el = points_[idxd].data_[4] - Wc.data_[4];
						diff += el*el;
					}

					if (diff < 1.0)
					{
						Mt += points_[idxd];
						pt_count++;
					}
				}
				idxd = slist[idxd];      // (most closed )previous point in the same bucket
			}
		}
		if( pt_count > 0 ){
			Mt /= (float)pt_count;
			Mt -= Wc;
		}
		MtAbs = 0.0f;
		for(j = 0; j < d_; j++)
			MtAbs += Mt.data_[j]*Mt.data_[j];

		iterationCount = 1;
		while((MtAbs >= EPSILON)&&(iterationCount < MAX_ITR))
		{
			pt_count = 0;
			Wc += Mt;	// Shift window location
			Mt.zero();

			// find bucket of yk
			cBuck1 = (int) Wc.data_[0] + 1;
			cBuck2 = (int) Wc.data_[1] + 1;
			cBuck3 = (int)(Wc.data_[2] - sMins) + 1;
			cBuck = cBuck1 + nBuck1*(cBuck2 + nBuck2*cBuck3);
			for (j = 0; j < 27; j++)
			{
				idxd = buckets[cBuck+bucNeigh[j]];
				// list parse, crt point is cHeadList
				while (idxd >= 0)
				{
					// determine if inside search window
					el = points_[idxd].data_[0]-Wc.data_[0];
					diff = el*el;
					el = points_[idxd].data_[1]-Wc.data_[1];
					diff += el*el;

					if (diff < 1.0)
					{
						el = points_[idxd].data_[2]-Wc.data_[2];
						if (Wc.data_[2] > hiLTr)
							diff = 4*el*el;
						else
							diff = el*el;

						if (d_ == 5)
						{
							el = points_[idxd].data_[3]-Wc.data_[3];
							diff += el*el;
							el = points_[idxd].data_[4]-Wc.data_[4];
							diff += el*el;
						}

						if (diff < 1.0)
						{
							Mt += points_[idxd];
							pt_count++;
						}
					}
					idxd = slist[idxd];
				}
			}
			if( pt_count > 0 ){
				Mt /= (float)pt_count;
				Mt -= Wc;
			}
			MtAbs = (Mt.data_[0]*Mt.data_[0]+Mt.data_[1]*Mt.data_[1]); //*sigmas*sigmas;
			if (d_==5)
				MtAbs += (Mt.data_[2]*Mt.data_[2]+Mt.data_[3]*Mt.data_[3]+Mt.data_[4]*Mt.data_[4]); //*sigmaR*sigmaR;
			else
				MtAbs += Mt.data_[2]*Mt.data_[2]; //*sigmaR*sigmaR;

			// Increment interation count
			iterationCount++;
		}

		Wc += Mt;
		//store the converge points
		for(j = 0; j < d_; j++)
			convpts_[i*d_+j] = Wc.data_[j];
	}

	delete []buckets;
	delete []slist;
	delete []wc;
	delete []mt;
}

